package com.tianji.aigc.memory;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.constant.ToolConstant;
import com.tianji.aigc.entity.RedisMessage;
import com.tianji.aigc.tools.result.ToolResultHolder;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * @version: 1.0
 * @Author: ljy
 * @description: AI机器人会话记忆类，用于保存用户会话信息
 * @date: 2025-09-02 19:58
 */
@Component
public class RedisChatMemory implements ChatMemory {

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 保存会话对应的消息到Redis中
     *
     * @param conversationId 对话ID
     * @param messages       对话列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        if (CollUtil.isEmpty(messages)) {
            return;
        }
        //获取redisKey，组成结构：CHAT:用户id:sessionId
        String key = this.getKey(conversationId);
        //将Message消息转为JSON字符串
        List<String> messageList = messages.stream().map(message -> {
            //拷贝Message对象为自定义的RedisMessage对象
            RedisMessage redisMessage = BeanUtil.toBean(message, RedisMessage.class);
            //设置消息内容字符串，否则为空
            redisMessage.setTextContent(message.getText());
            //设置Tool相关属性的值
            if (message instanceof AssistantMessage assistantMessage) {
                //将ToolResultHolder中存储的param数据存到redis缓存中
                //1. 得到messageId
                String messageId = (String)message.getMetadata().get("id");
                //2. 得到requestId
                String requestId = (String)ToolResultHolder.get(messageId, ToolConstant.REQUEST_ID);
                //3. 根据requestId得到param
                if (StrUtil.isNotBlank(requestId)) {
                    Map<String, Object> paramResult = ToolResultHolder.get(requestId);
                    if (CollUtil.isNotEmpty(paramResult)) {
                        redisMessage.setParams(paramResult);
                    }
                }
                redisMessage.setToolCalls(assistantMessage.getToolCalls());
            } else if (message instanceof ToolResponseMessage toolResponseMessage) {
                redisMessage.setToolResponses(toolResponseMessage.getResponses());
            }
            return JSONUtil.toJsonStr(redisMessage);
        }).toList();
        //将消息JSON字符串列表添加到Redis的List中
        redisTemplate.opsForList().leftPushAll(key, messageList);
    }

    /**
     * 从Redis中查询某会话下的指定条数的消息列表
     *
     * @param conversationId 对话ID
     * @param lastN          消息数目
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        if (lastN <= 0) {
            return List.of();
        }
        //获取redisKey，组成结构：CHAT:用户id:sessionId
        String key = this.getKey(conversationId);
        //从Redis的List中获取消息字符串列表
        List<String> messageStrList = redisTemplate.opsForList().range(key, 0, lastN);
        if (CollUtil.isEmpty(messageStrList)) {
            return List.of();
        }
        //将消息JSON字符串转为Message
        List<Message> messageList = messageStrList.stream().map(m -> {
            //将消息JSON字符串转为RedisMessage对象
            RedisMessage redisMessage = JSONUtil.toBean(m, RedisMessage.class);
            Message message = null;
            //根据消息类型封装具体消息对象
            if (redisMessage.getMessageType().equals(MessageType.USER)) {
                message = new UserMessage(redisMessage.getMessageType(), redisMessage.getTextContent(), redisMessage.getMedia(), redisMessage.getMetadata());
            } else if (redisMessage.getMessageType().equals(MessageType.ASSISTANT)) {
                message = new AssistantMessage(redisMessage.getTextContent(), redisMessage.getMetadata(), redisMessage.getToolCalls(), redisMessage.getMedia());
            } else if (redisMessage.getMessageType().equals(MessageType.SYSTEM)) {
                message = new SystemMessage(redisMessage.getTextContent());
            } else if (redisMessage.getMessageType().equals(MessageType.TOOL)) {
                message = new ToolResponseMessage(redisMessage.getToolResponses(), redisMessage.getMetadata());
            }
            return message;
        }).toList();
        return messageList;
    }

    /**
     * 从Redis中删除某会话下的全部消息
     *
     * @param conversationId 对话ID
     */
    @Override
    public void clear(String conversationId) {
        String key = this.getKey(conversationId);
        redisTemplate.delete(key);
    }

    /**
     * 获取redis中的key
     *
     * @param conversationId
     * @return
     */
    private String getKey(String conversationId) {
        return "CHAT:" + conversationId;
    }
}
